import os
import datetime
import pandas as pd
import numpy as np
from tqdm import tqdm
from copy import deepcopy
from tools.utils import loc_collection, loc_zq_collection, load_config
from tools.utils import loc_db
from tools.utils import get_ts_code_to_chinese_name


def get_pred_daily(ts_code, coll):
    y_name = f"{ts_code.upper()}_close"
    data = list(coll.find({'y_name': y_name}))
    return pd.DataFrame(data)


def get_pred_weekly(pred_code):
    dir_name = "data/pred_data/2022-01-11"
    print(dir_name)
    for pth in os.listdir(dir_name):
        data = pd.read_excel(f"{dir_name}/{pth}", index_col=0)
        columns = data.columns.to_list()
        print(columns)
        if pred_code in columns:
            idx = columns.index(pred_code)
            fit_data = data.loc[:, [columns[idx + 1]]]
            fit_data = fit_data.dropna(axis=0, how='all')
            return fit_data
    return pd.DataFrame()


def get_pred_by_history(dic: dict):
    cfg = load_config()
    if dic['freq'] == '日':
        db = loc_db('stock')
        latest_t = sorted(db.list_collection_names())[-1]
        print('latest t', latest_t)
        coll_ = loc_collection('stock', latest_t)
        d = get_pred_daily(dic['ts_code'], coll_)
        if len(d) < 1:
            return []
        print(dic['ts_code'], 'ts_code')
        t = d.iloc[0, :]['t']
        thd = 0.7
        d7 = d[d['acc'] > thd]
        while len(d7) < 1:
            thd -= 0.05
            d7 = d[d['acc'] > thd]
        acc_li = [v['acc'] for k, v in d7.iterrows()]
        print('np.max', acc_li)
        max_acc = np.max(acc_li) - 0.00001
        max_pred = {v['_id']: v['y_pred'] for k, v in d7.iterrows() if v['acc'] > max_acc}
        max_df = pd.DataFrame(max_pred)
        max_df = pd.DataFrame(max_df.mean(axis=1))
        max_df.columns = dic['columns'][:1]
        max_df.index = [str(s)[:10] for s in t]
        max_df = max_df.iloc[-9:, :]
    elif dic['freq'] == '周':
        map_ = get_ts_code_to_chinese_name(cfg)
        pred_code = map_[dic['ts_code']]
        max_df = get_pred_weekly(pred_code)
        if len(max_df) < 1:
            return []
        max_df.index = [str(s)[:10] for s in max_df.index]
        max_df.columns = dic['columns'][:1]
        thd = 0.7
    else:
        thd = 0
        max_df = pd.DataFrame()
    new_dic = deepcopy(dic)
    new_dic['columns'] = max_df.columns.to_list()
    new_dic['t'] = max_df.index.to_list()
    new_dic['value'] = max_df.values.tolist()
    new_dic['pred'] = True
    new_dic['acc'] = thd
    del new_dic['_id']
    return [new_dic]


def upload_informer_pred():
    cfg = load_config()
    db_name = "informer_weekly"
    td = str(datetime.date.today())
    today = 'goods'
    coll_data = loc_zq_collection('to_it', today)
    data_exists = list(coll_data.find({'freq': '周'}))
    db = loc_db(db_name, cfg)
    name = sorted(db.list_collection_names())[0]
    coll_pred = loc_collection(db_name, name)
    all_li = []
    p_bar = tqdm(data_exists, ncols=98)
    for dic in p_bar:
        find_data = list(coll_pred.find({'str_name': dic['columns'][0]}))
        if len(find_data) > 0:
            door_acc = 0.7
            while len([s for s in find_data if s['acc'] > door_acc]) == 0:
                door_acc -= 0.05
            acc_li = [v['acc'] for v in find_data if v['acc'] > door_acc]
            t_all = find_data[0]['t']
            t_fit = [tt for tt in t_all if tt > td]
            dic['t'] = find_data[0]['t'][-len(t_fit):]
            df = pd.concat([pd.Series(s['y_pred']) for s in find_data if s['acc'] > door_acc], axis=1)
            dic['acc'] = np.mean(acc_li)
            dic['value'] = df.mean(axis=1).to_list()[-len(t_fit):]
            dic['value_up'] = df.max(axis=1).to_list()[-len(t_fit):]
            dic['value_down'] = df.min(axis=1).to_list()[-len(t_fit):]
            dic['columns'] = [dic['columns'][0]]
            dic['pred'] = True
            del dic['_id']
            all_li.append(dic)
            p_bar.set_postfix({'name': dic['columns'][0]})
    coll_data.insert_many(all_li)


if __name__ == "__main__":
    upload_informer_pred()
